
#pragma once 

#include <vtkm/cont/ArrayHandle.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>
#include <vtkm/cont/arg/ControlSignatureTagBase.h>
#include <vtkm/cont/ConvertNumComponentsToOffsets.h>

using vtkm::worklet::WorkletMapField;
using vtkm::worklet::WorkletVisitCellsWithPoints;

#include <vtkm/cont/Invoker.h>
#include <vtkm/cont/CellSetExplicit.h>


#include "TypeCheck.h"
#include "Transport.h"
#include "Fetch.h"

#include "libmesh/elem.h"
#include "libmesh/mesh.h"
#include "libmesh/equation_systems.h"
#include <iostream>
#include <sstream>
#include <algorithm>
#include <math.h>

#include "libmesh/libmesh.h"
#include "libmesh/mesh.h"
#include "libmesh/mesh_generation.h"
#include "libmesh/exodusII_io.h"
#include "libmesh/linear_implicit_system.h"
#include "libmesh/equation_systems.h"

#include "libmesh/fe.h"

#include "libmesh/quadrature.h"


#include "libmesh/sparse_matrix.h"
#include "libmesh/numeric_vector.h"
#include "libmesh/dense_matrix.h"
#include "libmesh/dense_vector.h"
#include "libmesh/dof_map.h"

// To impose Dirichlet boundary conditions
#include "libmesh/dirichlet_boundaries.h"
#include "libmesh/analytic_function.h"

// The definition of a geometric element
#include "libmesh/elem.h"
#include "libmesh/enum_solver_package.h"
#include "libmesh/enum_quadrature_type.h"

#include "libmesh/enum_solver_type.h"
#include "libmesh/enum_preconditioner_type.h"
#include "libmesh/linear_solver.h"

using SystemBase = EbESystem;
#include "EbESystem.h"


struct AssemblyWorklet : WorkletVisitCellsWithPoints
{
	struct MeshIn : vtkm::cont::arg::ControlSignatureTagBase
	{
		using TypeCheckTag = TypeCheckTagMesh;
		using TransportTag = TransportTagMeshIn;
		using FetchTag = FetchTagMeshIn;
	};

	using ControlSignature = void(CellSetIn cellset);
	using ExecutionSignature = void(WorkIndex);
	
	libMesh::EquationSystems& es;

	AssemblyWorklet(libMesh::EquationSystems& es):es(es){}

	void operator()(vtkm::Id index) const
	{
		const MeshBase& mesh = es.get_mesh();
		const unsigned int dim = mesh.mesh_dimension();

		auto& system = es.get_system<SystemBase>("Poisson");
		const DofMap& dof_map = system.get_dof_map();

		FEType fe_type = dof_map.variable_type(0);
		std::unique_ptr<FEBase> fe(FEBase::build(dim, fe_type));
		std::unique_ptr<QBase> qrule(QBase::build(QuadratureType::QGAUSS, dim, THIRD));
		fe->attach_quadrature_rule(qrule.get());
		const std::vector<Real>& JxW = fe->get_JxW();
		const std::vector<std::vector<Real>>& phi = fe->get_phi();
		const std::vector<std::vector<RealGradient>>& dphi = fe->get_dphi();

		DenseMatrix<Number> Ke;
		DenseVector<Number> Fe;
		std::vector<dof_id_type> dof_indices;
		auto elem = mesh.elem_ptr(index);
		dof_map.dof_indices(elem, dof_indices);

		const unsigned int n_dofs =
			cast_int<unsigned int>(dof_indices.size());

		fe->reinit(elem);

		libmesh_assert_equal_to(n_dofs, phi.size());

		Ke.resize(n_dofs, n_dofs);

		Fe.resize(n_dofs);
		for (unsigned int qp = 0; qp < qrule->n_points(); qp++)
		{
			// Add the matrix contribution
			for (unsigned int i = 0; i != n_dofs; i++)
				for (unsigned int j = 0; j != n_dofs; j++)
					Ke(i, j) += JxW[qp] * (dphi[i][qp] * dphi[j][qp]);

			const Real fxy = 0;// -(uxx + uyy + ((dim == 2) ? 0. : uzz));

			for (unsigned int i = 0; i != n_dofs; i++)
				Fe(i) += JxW[qp] * fxy * phi[i][qp];
		}

		dof_map.heterogenously_constrain_element_matrix_and_vector(Ke, Fe, dof_indices);
		system.matrix->add_matrix(Ke, dof_indices);
		system.rhs->add_vector(Fe, dof_indices);

		system.add_matrix(Ke);
		system.add_rhs(Fe);
	}
};

vtkm::cont::CellSetExplicit<> BuildCellSet(const MeshBase& mesh)
{
	vtkm::cont::CellSetExplicit<> cs;
	auto n_pts = static_cast<vtkm::Id>(mesh.n_nodes());
	auto n_els = static_cast<vtkm::Id>(mesh.n_elem());
	std::vector<vtkm::UInt8> shapes;
	std::vector<vtkm::Id> connectivity;
	//std::vector<vtkm::Id> offsets;
	std::map<libMesh::ElemType, vtkm::CellShapeIdEnum>  elem_type;
	elem_type.emplace(libMesh::ElemType::QUAD4, vtkm::CELL_SHAPE_QUAD);
	elem_type.emplace(libMesh::ElemType::TRI3, vtkm::CELL_SHAPE_TRIANGLE);
	std::vector<vtkm::Id> n_nodes;

	//offsets.push_back(0);
	for (auto elem : mesh.active_local_element_ptr_range())
	{
		shapes.push_back(elem_type[elem->type()]);
		//offsets.push_back(elem->n_nodes());
		for (size_t i = 0; i < elem->n_nodes(); i++)
		{
			connectivity.push_back(elem->node_id(i));
		}
	}
	auto offsets = vtkm::cont::ConvertNumComponentsToOffsets(vtkm::cont::make_ArrayHandle(n_nodes, vtkm::CopyFlag::Off));

	cs.Fill(n_pts, 
		vtkm::cont::make_ArrayHandle(shapes, vtkm::CopyFlag::Off),
		vtkm::cont::make_ArrayHandle(connectivity, vtkm::CopyFlag::Off),
		offsets
		);
	return cs;
}

void Run(libMesh::EquationSystems &es)
{
	const MeshBase& mesh = es.get_mesh();

	vtkm::cont::ArrayHandle<libMesh::Elem*> elements;
	elements.Allocate(mesh.n_elem());
	for (auto elem : mesh.active_local_element_ptr_range())
	{
		elements.WritePortal().Set(elem->id(), elem);
	}

	auto cellset = BuildCellSet(mesh);
	vtkm::cont::Invoker invoker;
	invoker(AssemblyWorklet{es}, cellset);
}

